package lnet

import (
	"sync"
	"sync/atomic"
)

type KEY interface{}

type Channel struct {
	mutex sync.RWMutex
	sesss map[KEY]*Session
	State atomic.Value
}

func NewChannel() *Channel {
	return &Channel{
		sesss: make(map[KEY]*Session),
	}
}

func (c *Channel) Len() int {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	return len(c.sesss)
}

func (c *Channel) Fetch(callback func(*Session)) {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	for _, sess := range c.sesss {
		callback(sess)
	}
}

func (c *Channel) Get(key KEY) *Session {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	sess, _ := c.sesss[key]
	return sess
}

func (c *Channel) Put(key KEY, sess *Session) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if sess, exists := c.sesss[key]; exists {
		c.remove(key, sess)
	}
	sess.AddCloseCallback(c, key, func() {
		c.Remove(key)
	})
	c.sesss[key] = sess
}

func (c *Channel) remove(key KEY, sess *Session) {
	sess.RemoveCloseCallback(c, key)
	delete(c.sesss, key)
}

func (c *Channel) Remove(key KEY) bool {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	sess, exists := c.sesss[key]
	if exists {
		c.remove(key, sess)
	}
	return exists
}

func (c *Channel) FetchAndRemove(callback func(*Session)) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	for key, sess := range c.sesss {
		sess.RemoveCloseCallback(c, key)
		delete(c.sesss, key)
		callback(sess)
	}
}

func (c *Channel) Close() {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	for key, sess := range c.sesss {
		c.remove(key, sess)
	}
}
